#include "ST_EstimateFunction.h"

#include <memory.h>

#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_multifit.h>


/* 3 by 3 matrix inverse */
void Inv33(double *src, double *dst)
{
	double p1 = src[0] * (src[4] * src[8] - src[5] * src[7]);
	double p2 = src[1] * (src[3] * src[8] - src[5] * src[6]);
	double p3 = src[2] * (src[3] * src[7] - src[4] * src[6]);

	double det = p1 - p2 + p3;

	double t[9];
	t[0] = src[0]; t[1] = src[3]; t[2] = src[6];
	t[3] = src[1]; t[4] = src[4]; t[5] = src[7];
	t[6] = src[2]; t[7] = src[5]; t[8] = src[8];

	dst[0] = (t[4] * t[8] - t[5] * t[7]) / det * 1.0;
	dst[1] = (t[3] * t[8] - t[5] * t[6]) / det * -1.0;
	dst[2] = (t[3] * t[7] - t[4] * t[6]) / det * 1.0;

	dst[3] = (t[1] * t[8] - t[2] * t[7]) / det * -1.0;
	dst[4] = (t[0] * t[8] - t[2] * t[6]) / det * 1.0;
	dst[5] = (t[0] * t[7] - t[1] * t[6]) / det * -1.0;

	dst[6] = (t[1] * t[5] - t[2] * t[4]) / det * 1.0;
	dst[7] = (t[0] * t[5] - t[2] * t[3]) / det * -1.0;
	dst[8] = (t[0] * t[4] - t[1] * t[3]) / det * 1.0;
}

void LeastSquareMY(float *A, float *x, float *y, const int n, const int m)
{
	if (m != 3)
		return;

	double *At = new double[n*m];
	for (int i = 0; i < n; ++i) {
		for (int j = 0; j < m; ++j) {
			At[j*n+i] = A[i*m+j];
		}
	}

	double *AA = new double[m*m];
	for (int i = 0; i < m; ++i) {
		for (int j = 0; j < m; ++j) {
			AA[i*m+j] = 0.0;
			for (int k = 0; k < n; ++k) {
				AA[i*m+j] += At[i*n+k] * A[k*m+j];
			}
		}
	}

	double *invAA = new double[m*m];
	Inv33(AA, invAA);

	double *AAA = new double[m*n];
	for (int i = 0; i < m; ++i) {
		for (int j = 0; j < n; ++j) {
			AAA[i*n+j] = 0.0f;
			for (int k = 0; k < m; ++k) {
				AAA[i*n+j] += invAA[i*m+k] * At[k*n+j];
			}
		}
	}

	for (int i = 0; i < m; ++i) {
		x[i] = 0.0f;
		for (int j = 0; j < n; ++j) {
			x[i] += (float)(AAA[i*n+j]) * y[j];
		}
	}

	delete[] At;
	delete[] AA;
	delete[] invAA;
	delete[] AAA;
}

void LeastSquareGSL(float *A, float *x, float *y, const int n, const int m)
{
    /* allocate */
    gsl_matrix *gslA, *gslCov;
    gsl_vector *gslx, *gsly;
    double chisq = 0.0;

    gslA = gsl_matrix_alloc(n, m);
    gslx = gsl_vector_alloc(m);
    gsly = gsl_vector_alloc(n);
    gslCov = gsl_matrix_alloc(m, m);

    /* set */
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < m; ++j) {
            gsl_matrix_set(gslA, i, j, A[i*m+j]);
        }
        gsl_vector_set(gsly, i, y[i]);
    }

    /* run */
    gsl_multifit_linear_workspace *work = gsl_multifit_linear_alloc(n, m);
    gsl_multifit_linear(gslA, gsly, gslx, gslCov, &chisq, work);

    /* get */
    for (int j = 0; j < m; ++j) {
        x[j] = gsl_vector_get(gslx, j);
    }

    /* debug */
    /*printf("Matrix A:\n");
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < m; ++j) {
            printf("%f ", A[i*m+j]);
        }
        printf("\n");
    }*/

    /* free */
    gsl_multifit_linear_free(work);
    gsl_matrix_free(gslA);
    gsl_matrix_free(gslCov);
    gsl_vector_free(gslx);
    gsl_vector_free(gsly);
}

void EstimateFunction(CPoint3F *F, CPoint3F *X,
        const int m, const int n, float *function)
{
    /* matrix A is same for all parameters */
    float *A = new float[m*3];
    for (int k = 0; k < m; ++k) {
        A[k*3+0] = X[k].m_x;
        A[k*3+1] = X[k].m_y;
        A[k*3+2] = X[k].m_z;
    }

    for (int i = 0; i < n; ++i) {
        /* estimate the function for each parameter */
        for (int j = 0; j < 3; ++j) {
            /* estimat the function for each compoment of each parameter */

            float *x = new float[3];
            float *y = new float[m];

            for (int k = 0; k < m; ++k) {
                if (j == 0) {
                    y[k] = F[k*n+i].m_x;
                } else if (j == 1) {
                    y[k] = F[k*n+i].m_y;
                } else if (j == 2) {
                    y[k] = F[k*n+i].m_z;
                }
            }
            memset(x, 0, sizeof(float)*3);


            /* using least square to estimate the function */
            //LeastSquareGSL(A, x, y, m, 3);
			LeastSquareMY(A, x, y, m, 3);

            /* get the result */
            function[i*9+j*3+0] = x[0];
            function[i*9+j*3+1] = x[1];
            function[i*9+j*3+2] = x[2];

            delete[] x; x = NULL;
            delete[] y; y = NULL;
        }
    }

    delete[] A; A = NULL;
}
